#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from typing import Annotated
from fastapi import APIRouter, Depends, Path, Query, Request
from fastapi.encoders import jsonable_encoder
from fastapi.responses import StreamingResponse, JSONResponse
import json

from starlette.websockets import WebSocketState

from app.crud.crud_conservation import conservation_dao
from app.schema import (
    CreateAgentParam,
    AddAgentPluginParam,
    AddAgentKnowledgeBaseParam,
    UpdateAgentParam,
    QueryAgentParam, GetAgentList, GetAgentDetail,
)
from app.schema.conservation_schema import CreateConservationParam, UpdateConservationParam
from app.service.agent_service import agent_service
from app.service.conservation_service import conservation_service
from common.pagination import DependsPagination, paging_data
from common.response.response_schema import ResponseModel, response_base
from common.security.jwt import DependsJwtAuth
from common.security.permission import RequestPermission
from database.db_mysql import CurrentSession, async_db_session
from utils.serializers import select_as_dict
from fastapi import WebSocket, WebSocketDisconnect
import logging

logger = logging.getLogger(__name__)
router = APIRouter()


@router.get('/all', summary='查看所有智能体', dependencies=[DependsJwtAuth])
async def get_all_agents(
        request: Request,
        description: Annotated[str | None, Query()] = None,
        name: Annotated[str | None, Query()] = None,
) -> ResponseModel:
    # 获取所有会话
    agents = await agent_service.get_all(user_uuid=request.user.uuid, description=description, name=name)

    # 转换为响应所需的格式
    data = [GetAgentList(**select_as_dict(agent)) for agent in agents]
    return response_base.success(data=data)


@router.get(
    '/all/public',
    summary='分页获取所有智能体',
    dependencies=[
        DependsJwtAuth,
        DependsPagination,
    ],
)
async def get_pagination_public_agents(
        request: Request,
        db: CurrentSession,
        name: Annotated[str | None, Query()] = None,
        description: Annotated[str | None, Query()] = None,
):
    agent_select = await agent_service.get_select(name=name, description=description, status=2)
    page_data = await paging_data(db, agent_select, GetAgentList)
    return response_base.success(data=page_data)


@router.post('', summary='添加智能体', dependencies=[DependsJwtAuth])
async def add_agent(request: Request, obj: CreateAgentParam) -> ResponseModel:
    obj.user_uuid = request.user.uuid
    agent = await agent_service.add(obj=obj)
    await conservation_service.add(request=request, obj=CreateConservationParam(
        user_uuid=request.user.uuid,
        agent_uuid=agent.uuid,
    ))
    return response_base.success()


@router.put('/{agent_uuid}', summary='更新智能体信息', dependencies=[DependsJwtAuth])
async def update_agent(request: Request, agent_uuid: Annotated[str, Path(...)], obj: UpdateAgentParam) -> ResponseModel:
    count = await agent_service.update(request=request, agent_uuid=agent_uuid, obj=obj)
    if count > 0:
        return response_base.success()
    return response_base.fail()


@router.put('/{agent_uuid}/status', summary='修改智能体状态', dependencies=[DependsJwtAuth])
async def update_agent_status(request: Request, agent_uuid: Annotated[str, Path(...)],
                              status: Annotated[bool, Query()]) -> ResponseModel:
    count = await agent_service.update_status(request=request, agent_uuid=agent_uuid, status=status)
    if count > 0:
        return response_base.success()
    return response_base.fail()


@router.post('/{agent_uuid}/plugin', summary='给智能体添加插件', dependencies=[DependsJwtAuth])
async def reset_plugin_to_agent(request: Request, agent_uuid: Annotated[str, Path(...)],
                                plugin_param: AddAgentPluginParam) -> ResponseModel:
    await agent_service.add_plugin(request=request, agent_uuid=agent_uuid, plugin_param=plugin_param)
    return response_base.success()


@router.post('/{agent_uuid}/knowledge_base', summary='给智能体添加知识库', dependencies=[DependsJwtAuth])
async def reset_knowledge_base_to_agent(request: Request, agent_uuid: Annotated[str, Path(...)],
                                        knowledge_base_param: AddAgentKnowledgeBaseParam) -> ResponseModel:
    await agent_service.add_knowledge_base(request=request, agent_uuid=agent_uuid,
                                           knowledge_base_param=knowledge_base_param)
    return response_base.success()


@router.get('/{agent_uuid}', summary='查看智能体信息', dependencies=[DependsJwtAuth])
async def get_agent_info(request: Request, agent_uuid: Annotated[str, Path(...)]) -> ResponseModel:
    current_agent = await agent_service.get_with_relation(request=request, agent_uuid=agent_uuid)
    data = GetAgentDetail(**select_as_dict(current_agent), operator_uuid=request.user.uuid)
    return response_base.success(data=data)


@router.get(
    '',
    summary='分页获取所有智能体',
    dependencies=[
        DependsJwtAuth,
        DependsPagination,
    ],
)
async def get_pagination_agents(
        request: Request,
        db: CurrentSession,
        name: Annotated[str | None, Query()] = None,
        description: Annotated[str | None, Query()] = None,
        status: Annotated[int | None, Query()] = None,
):
    agent_select = await agent_service.get_select(user_uuid=request.user.uuid, name=name, description=description,
                                                  status=status)
    page_data = await paging_data(db, agent_select, GetAgentList)
    return response_base.success(data=page_data)


@router.delete(
    path='/{agent_uuid}',
    summary='删除智能体',
    description='删除后智能体将从数据库中删除',
    dependencies=[DependsJwtAuth, Depends(RequestPermission('sys:agent:del'))],
)
async def delete_agent(request: Request, agent_uuid: Annotated[str, Path(...)]) -> ResponseModel:
    count = await agent_service.delete(request=request, agent_uuid=agent_uuid)  # no need for request here
    if count > 0:
        return response_base.success()
    return response_base.fail()


@router.get("/generate_spl_form/{agent_uuid}")
async def require_2_spl_form(request: Request, agent_uuid: Annotated[str, Path(...)]):
    return StreamingResponse(agent_service.generate_spl_form(request=request, agent_uuid=agent_uuid),
                             media_type="text/plain")


@router.get(
    "/generate_spl_chain/{agent_uuid}",
    dependencies=[DependsJwtAuth],
)
async def require_2_spl_chain(request: Request, agent_uuid: Annotated[str, Path(...)]):
    return StreamingResponse(agent_service.generate_spl_chain(request=request, agent_uuid=agent_uuid),
                             media_type="text/plain")


@router.post(
    "/generate_answer/{conservation_uuid}",
    dependencies=[DependsJwtAuth],
)
async def run_spl_chain(request: Request, conservation_uuid: Annotated[str, Path(...)], query: QueryAgentParam):
    return StreamingResponse(agent_service.generate_answer(conservation_uuid=conservation_uuid, query=query.message),
                             media_type="text/plain")


@router.websocket("/ws/generate_answer/{conservation_uuid}")
async def run_spl_chain_socket(websocket: WebSocket):
    # 1. 正确获取路径参数
    conservation_uuid = websocket.path_params.get("conservation_uuid")
    if not conservation_uuid:
        await websocket.close(code=4001, reason="Missing conservation_uuid in path")
        return

    await websocket.accept()
    logger.info(f"WebSocket connected for conservation: {conservation_uuid}")
    llm_res = ''
    try:
        # 2. 接收初始查询消息
        data = await websocket.receive_text()
        query = json.loads(data).get('message')
        logger.debug(f"Received query for {conservation_uuid}: {query}")

        # 3. 初始化 chat_history
        async with async_db_session() as db:
            conservation = await conservation_dao.get_with_relation(db=db, conservation_uuid=conservation_uuid)
            history = json.loads(conservation.chat_history)
            history.append({"role": "user", "contents": query})

        # 4. 流式生成响应
        response_stream = agent_service.generate_answer(
            conservation_uuid=conservation_uuid,
            query=query
        )
        # 5. 流式发送响应并记录 chat_history
        async for part in response_stream:
            if websocket.client_state == WebSocketState.DISCONNECTED:
                logger.warning(f"Client disconnected early for {conservation_uuid}")
                break
            llm_res = json.loads(part)
            await websocket.send_text(part)
            logger.debug(f"Sent part for {conservation_uuid}: {part}")

        history.append({"role": 'system', "contents": llm_res})
        # 保存当前的 chat_history
        async with async_db_session() as db:
            await conservation_dao.update(
                db,
                conservation_uuid=conservation_uuid,
                obj=UpdateConservationParam(chat_history=history)
            )
            await db.commit()

        # 6. 正常完成关闭
        logger.info(f"Completed streaming for {conservation_uuid}")
        await websocket.close(code=1000, reason="Stream completed")

    except WebSocketDisconnect as e:
        # 客户端主动断开
        logger.info(f"Client disconnected for {conservation_uuid}: code={e.code}, reason={e.reason}")
        history.append({"role": 'system', "contents": llm_res})
        # 保存当前的 chat_history
        async with async_db_session() as db:
            await conservation_dao.update(
                db,
                conservation_uuid=conservation_uuid,
                obj=UpdateConservationParam(chat_history=history)
            )
            await db.commit()
        logger.info(f"Saved chat_history for {conservation_uuid} after client disconnect")

    except Exception as e:
        # 异常处理
        logger.error(f"Error processing {conservation_uuid}: {str(e)}", exc_info=True)
        history.append({"role": 'system', "contents": llm_res})
        # 保存当前的 chat_history
        async with async_db_session() as db:
            await conservation_dao.update(
                db,
                conservation_uuid=conservation_uuid,
                obj=UpdateConservationParam(chat_history=history)
            )
            await db.commit()
        logger.info(f"Saved chat_history for {conservation_uuid} after server error")

        # 尝试关闭 WebSocket
        try:
            await websocket.close(code=1011, reason=f"Server error: {str(e)}")
        except:
            pass


@router.post(
    "/download/{agent_uuid}",
    dependencies=[DependsJwtAuth],
)
async def download_agent_endpoint(request: Request, agent_uuid: Annotated[str, Path(...)]):
    download_data = await agent_service.download_agent(request=request, agent_uuid=agent_uuid)
    json_data = jsonable_encoder(download_data)
    return JSONResponse(content=json_data, media_type="application/json", headers={
        "Content-Disposition": f"attachment; filename=agent_data.json"
    })
